package org.ala.tree.bat;

import java.util.Comparator;
import java.util.stream.IntStream;

import org.ala.tree.AbstractTree;

/**
 * B+树
 * <p>	与B树不同，所有数据都存储在叶子节点上
 * <p>	叶子节点会有一层链表指针，指向后继或前驱
 * <p>	m阶b+树，每个节点中最多有m个关键字，最少有m//2个关键字
 * <p>		每个节点是其子节点关键字的最大值
 * <p>		与b树不同，b+树的关键字个数就是阶数
 * <p>	b+树节点上的每个关键字，代表其指向的子节点的关键字最大值
 * <p>	
 * <p>	这里仅仅是个演示，关键字和节点列表是用数组存的（在分裂和合并时会有大量的移动操作）。
 *
 * @author ala
 * @date 2025年5月20日
 */
public class BATree<K,V> extends AbstractTree<K, V> {
	/**
	 * 	根节点，叶子节点双链头尾针
	 */
	protected Node<K,V> ROOT, HEAD, TAIL;
	/**
	 * 阶数
	 */
	protected int m;
	/**
	 * 每层最小节点数
	 */
	protected int mn;
	/**
	 * 表示无穷大的k
	 */
	protected K MAX_K;
	
	
	public BATree(int m, K maxK, Comparator<K> comparator) {
		super(comparator);
		if (m < 3) {throw new RuntimeException("阶数最小是3");}
		this.m = m;
		mn = m >> 1;
		MAX_K = maxK;
		
		ROOT = buildLear();
		//	追加一个无穷大哨兵
		insertKV(ROOT, MAX_K, null);
		
		HEAD = buildNode();
		TAIL = buildNode();
		HEAD.next = ROOT; ROOT.next = TAIL;
		TAIL.prev = ROOT; ROOT.prev = HEAD;
	}

	@Override
	public boolean add(K k, V v) {
		if (_add(ROOT, k, v)) {size++; return true;}
		return false;
	}

	@Override
	public boolean remove(K k) {
		if (_remove(k)) {size--; return true;}
		return false;
	}

	@Override
	public boolean update(K k, V v) {
		SearchResult<K,V> sr = _search(ROOT, -1, k);
		if (!sr.find) {return false;}
		sr.node.vals[sr.ki] = v;
		return true;
	}

	@Override
	public V search(K k) {
		SearchResult<K,V> sr = _search(ROOT, -1, k);
		return sr.find ? sr.val : null;
	}
	
	
	/********************************************************************************************************
	 * 插入相关操作
	 ********************************************************************************************************/
	protected boolean _add(Node<K,V> r, K k, V v) {
		SearchResult<K, V> sr = _search(r, -1, k);
		//	不允许重复
		if (sr.find) {return false;}
		
		//	此时n必为叶子节点
		Node<K,V> n = sr.node;
		insertKV(n, k, v);
		tryUp(n);
		return true;
	}
	/**
	 * 	把键值对插入叶子节点中
	 */
	@SuppressWarnings("unchecked")
	protected void insertKV(Node<K,V> n, K k, V v) {
		int i = n.kp;
		for (int t = 0 ; t < n.kp ; t++) {
			int c = comparator.compare(k, (K) n.keys[t]);
			if (c <= 0) {i = t; break;}
		}
		for (int t = n.kp ; t > i ; t--) {
			n.keys[t] = n.keys[t-1];
			n.vals[t] = n.vals[t-1];
		}
		n.keys[i] = k;
		n.vals[i] = v;
		n.kp++;
	}
	/**
	 * 	检测处理上溢
	 */
	protected void tryUp(Node<K, V> n) {
		if (n.kp <= m) {return;}
		//	从i位置开始分裂
		int i = n.kp >> 1;
		Node<K,V> nn = split(n, i);
		
		//	将原节点的末尾上移到父节点，并且指向原节点
		Node<K,V> p = n.parent;
		if (p == null) {
			p = buildNode();
			ROOT = p;
			p.keys[0] = MAX_K;
			p.nodes[0] = nn;
			p.kp = 1;
			n.parent = p;
		}
		nn.parent = p;
		upKey(n, p, nn);
		
		//	检测父节点是否发生上溢
		tryUp(p);
	}
	/**
	 * 	从i位置切节点
	 *  <p>	[0,i)保留原节点，[i,kp)放入新节点
	 *  <p>	原同层链表关系也维护好
	 *  @return 新切出来的节点（右半段，左半段在原节点）
	 */
	protected Node<K,V> split(Node<K,V> n, int i) {
		Node<K,V> nn;
		//	如果切的是叶子节点
		if (n.lear) {
			nn = buildLear();
			for (int t = i ; t < n.kp ; t++) {
				nn.keys[nn.kp] = n.keys[t];
				nn.vals[nn.kp++] = n.vals[t];
			}
		} 
		//	如果切的是非叶子
		else {
			nn = buildNode();
			for (int t = i ; t < n.kp ; t++) {
				nn.keys[nn.kp] = n.keys[t];
				nn.nodes[nn.kp] = n.nodes[t];
				nn.nodes[nn.kp++].parent = nn;
			}
		}
		n.kp = i;
		//	在叶子节点处维护链表
		if (n.lear) {insertNodeAfter(n, nn);}
		return nn;
	}
	/**
	 * 	把n的末尾关键字上移到父节点处
	 * 	<p>	并且上移的关键字指向n
	 */
	@SuppressWarnings("unchecked")
	protected void upKey(Node<K,V> n, Node<K,V> p, Node<K,V> nn) {
		//	找到插入的位置
		int i = p.kp;
		K k = (K) n.keys[n.kp-1];
		for (int t = 0 ; t < p.kp ; t++) {
			int c = comparator.compare(k, (K)p.keys[t]);
			if (c <= 0) {
				i = t; break;
			}
		}
		
		//	父节点的keys和nodes从i开始整体右移一位
		p.nodes[i] = nn;
		for (int j = p.kp ; j > i ; j--) {
			p.keys[j] = p.keys[j-1];
			p.nodes[j] = p.nodes[j-1];
		}
		//	在位置i插入n的末尾key，并且指向n
		p.keys[i] = n.keys[n.kp-1];
		p.nodes[i] = n;
		p.kp++;
	}
	
	
	/********************************************************************************************************
	 * 	删除相关
	 ********************************************************************************************************/
	protected boolean _remove(K k) {
		SearchResult<K,V> sr = _search(ROOT, -1, k);
		if (!sr.find) {return false;}
		
		removeKV(sr.node, sr.ki);
		tryDown(sr.node, sr.ni);
		return true;
	}
	/**
	 * 	从叶子节点中删除kv
	 */
	protected void removeKV(Node<K,V> n, int i) {
		for (int t = i ; t < n.kp-1 ; t++) {
			n.keys[t] = n.keys[t+1];
			if (n.lear) {
				n.vals[t] = n.vals[t+1];
			} else {
				n.nodes[t] = n.nodes[t+1];
			}
		}
		n.kp--;
	}
	/**
	 * 	检测是否下溢
	 */
	protected void tryDown(Node<K,V> n, Integer ni) {
		if (n.parent == null || n.kp >= mn) {return;}
		
		Node<K,V> p = n.parent;
		//	找到其左兄弟或者右兄弟
		if (ni == null) {
			for (int i = 0 ; i < p.kp ; i++) { 
				if (p.nodes[i] == n) {ni = i; break;}
			}
		}
		//	左右兄弟必有一个不为空。否则n就应该是父节点
		Node<K,V> nl = ni > 0 ? p.nodes[ni-1] : null,
				  nr = ni < p.kp-1 ? p.nodes[ni+1] : null;
		if (nl != null && nr != null) {
			//	找多的兄弟借
			if (nl.kp >= mn && nr.kp >= mn) {
				if (nl.kp >= nr.kp) { borrow(p, n, ni, nl, true); } 
				else { borrow(p, n, ni, nr, false); }
			} 
			else if (nl.kp >= mn) { borrow(p, n, ni, nl, true); } 
			else if (nr.kp >= mn) { borrow(p, n, ni, nr, false); } 
			else {
				//	往少的兄弟合并
				if (nl.kp <= nr.kp) {merge(p, n, ni, nl, true);}
				else {merge(p, n, ni, nr, false);}
				//	因为父节点要删除一个节点，检测父节点是否下溢
				tryDown(p, null);
			}
		} else if (nl != null) {
			if (nl.kp >= mn) { borrow(p, n, ni, nl, true); }
			else {merge(p, n, ni, nl, true); tryDown(p, null);}
		} else {
			if (nr.kp >= mn) { borrow(p, n, ni, nr, false); }
			else {merge(p, n, ni, nr, false); tryDown(p, null);}
		}
	}
	/**
	 * 	从兄弟节点借一个关键字到本节点
	 *  <p>	因为父节点的节点数没有变化，所以不会引发连锁反应
	 */
	protected void borrow(Node<K,V> p, Node<K,V> n, int ni, Node<K,V> nb, boolean left) {
		if (left) {
			//	把左节点的末尾key，追加到本节点（左节点的末尾key，等于父节点中指向左节点的key）
			for (int i = n.kp ; i >= 0 ; i--) {
				n.keys[i] = n.keys[i-1];
				if (n.lear) {n.vals[i] = n.vals[i-1];}
				else {n.nodes[i] = n.nodes[i-1];}
			}
			n.keys[0] = nb.keys[nb.kp-1];
			if (n.lear) { n.vals[0] = nb.vals[nb.kp-1]; }
			else {n.nodes[0] = nb.nodes[nb.kp-1];}
			n.kp++;
			//	把左节点的末尾覆盖到父节点的ni处
			p.keys[ni] = nb.keys[nb.kp-1];
			nb.kp--;
		} else {
			//	把右节点的开头，追加到本节点
			n.keys[n.kp] = nb.keys[0];
			if (n.lear) {n.vals[n.kp] = nb.vals[0]; }
			else {n.nodes[n.kp] = nb.nodes[0];}
			//	把本节点的最大元素，上移到父节点ni处
			p.keys[ni] = n.keys[n.kp];
			n.kp++;
			for (int i = 1 ; i < nb.kp ; i++) {
				nb.keys[i-1] = nb.keys[i];
				if (nb.lear) {nb.vals[i-1] = nb.vals[i];}
				else {nb.nodes[i-1] = nb.nodes[i];}
			}
			nb.kp--;
		}
	}
	/**
	 * 	合并节点
	 */
	protected void merge(Node<K,V> p, Node<K,V> n, int ni, Node<K,V> nb, boolean left) {
		//	往左合，把n复制给nl的后半段。删掉父节点的ni-1关键字和节点
		if (left) {
			//	把n复制给nl的后半段
			for (int i = 0 ; i < n.kp ; i++) {
				nb.keys[nb.kp] = n.keys[i];
				if (n.lear) {
					nb.vals[nb.kp] = n.vals[i];
				} else {
					nb.nodes[nb.kp] = n.nodes[i];
				}
				nb.kp++;
			}
			//	删掉父节点的ni-1关键字和节点
			for (int i = ni-1 ; i < p.kp ; i++) {
				p.keys[i] = p.keys[i+1];
				p.nodes[i] = p.nodes[i+1];
			}
			p.kp--;
		} 
		//	往右合，把n复制给nb的前半段。删掉父节点的ni关键字和节点
		else {
			//	把n复制给nb的前半段（把nb复制给n的后半段，再把n和nb做值替换）
			for (int i = 0 ; i < nb.kp ; i++) {
				n.keys[n.kp] = nb.keys[i];
				if (n.lear) {
					n.vals[n.kp] = nb.vals[i];
				} else {
					n.nodes[n.kp] = nb.nodes[i];
				}
				n.kp++;
			}
			nb.keys = n.keys;
			nb.nodes = n.nodes;
			nb.vals = n.vals;
			//	删掉父节点的ni关键字和节点
			for (int i = ni ; i < p.kp ; i++) {
				p.keys[i] = p.keys[i+1];
				p.nodes[i] = p.nodes[i+1];
			}
			p.kp--;
		}
		
		//	删除n节点
		removeNode(n);
		n.clear();
	}
	
	
	/********************************************************************************************************
	 * 	搜索相关
	 ********************************************************************************************************/
	/**
	 * 搜索关键字，和其所在的叶子节点。如果搜不到，返回其应该被insert的叶子节点
	 */
	@SuppressWarnings("unchecked")
	protected SearchResult<K,V> _search(Node<K,V> r, int ni, K k) {
		for (int i = 0 ; i < r.kp ; i++) {
			int c = comparator.compare(k, (K)r.keys[i]);
			if (r.lear) {
				if (c == 0) {return SearchResult.build(k, (V)r.vals[i], i, r, ni);}
				else {return SearchResult.build(r, ni);}
			} else {
				if (c <= 0) {return _search(r.nodes[i], i, k);}
			}
		}
		return SearchResult.build(r, ni);
	}
	
	
	/********************************************************************************************************
	 * 叶子节点的链表相关操作
	 ********************************************************************************************************/
	protected void insertNodeBefore(Node<K,V> nxt, Node<K,V> n) {
		Node<K,V> prv = nxt.prev;
		prv.next = n; n.next = nxt;
		nxt.prev = n; n.prev = prv;
	}
	protected void insertNodeAfter(Node<K,V> prv, Node<K,V> n) {
		Node<K,V> nxt = prv.next;
		prv.next = n; n.next = nxt;
		nxt.prev = n; n.prev = prv;
	}
	protected void removeNode(Node<K,V> n) {
		Node<K,V> prv = n.prev, nxt = n.next;
		prv.next = nxt; nxt.prev = prv;
		n.clear();
	}
	
	
	@SuppressWarnings("unchecked")
	protected Node<K,V> buildLear() {
		Object[] keys = new Object[m+1];
		Object[] vals = new Object[m+1];
		Node<K,V>[] nodes = new Node[m];
		return new Node<K,V>(keys, nodes, vals, true);
	}
	@SuppressWarnings("unchecked")
	protected Node<K,V> buildNode() {
		Object[] keys = new Object[m+1];
		Node<K,V>[] nodes = new Node[m+1];
		return new Node<K,V>(keys, nodes, null, false);
	}
	
	
	@Override
	public String toString() {
		StringBuilder sbuf = new StringBuilder();
		toString(sbuf, ROOT, 0);
		return sbuf.toString();
	}
	protected void toString(StringBuilder sbuf, Node<K,V> r, int level) {
		StringBuilder t = new StringBuilder();
		IntStream.range(0, level).forEach(i -> t.append("\t"));
		for (int i = 0 ; i < r.kp ; i++) {
			sbuf.append(t).append(r.keys[i]).append("\n");
			if (!r.lear) {
				toString(sbuf, r.nodes[i], level+1);
			}
		}
	}


	/**
	 * 	搜索结果
	 */
	public static class SearchResult<K,V> {
		boolean find;
		K key;
		V val;
		int ki;
		Node<K,V> node;
		int ni;
		public static <K,V> SearchResult<K,V> build(K k, V v, int ki, Node<K,V> node, int ni) {
			SearchResult<K,V> sr = new SearchResult<>();
			sr.key = k;
			sr.val = v;
			sr.ki = ki;
			sr.node = node;
			sr.ni = ni;
			sr.find = true;
			return sr;
		}
		public static <K,V> SearchResult<K,V> build(Node<K,V> node, int ni) {
			SearchResult<K,V> sr = new SearchResult<>();
			sr.node = node;
			sr.ni = ni;
			sr.find = false;
			return sr;
		}
	}
	/**
	 * 	m阶b+树节点信息
	 * 	m阶b+树，每层有m-1个关键字，有m个子节点，每个关键字是其子节点中关键字的最大值
	 */
	public static class Node<K,V> {
		boolean lear;
		Object[] keys;
		Node<K,V>[] nodes;
		int kp = 0;
		Object[] vals;
		Node<K,V> prev, next, parent;
		public Node(Object[] keys, Node<K,V>[] nodes, Object[] vals, boolean lear) {
			this.keys = keys;
			this.nodes = nodes;
			this.vals = vals;
			this.lear = lear;
		}
		public void clear() {
			keys = null;
			vals = null;
			nodes = null;
		}
		public String toString() {
			StringBuilder sbuf = new StringBuilder("[");
			for (int i = 0 ; i < kp ; i++) {
				if (i > 0) {sbuf.append(",");}
				sbuf.append(keys[i]);
			}
			sbuf.append("]");
			return sbuf.toString();
		}
		
	}
}
